import os
import cv2

import logging
import ddddocr

''' ------------------- setting配置文件 ----------------------------------'''
# 两个不一样的ddddocr配置, 可以选用一个是ddddocr第一套方案, 一个使用ddddocr第二套方案
# 电脑没有配置GPU的话使用把GPU参数改为False
ocr1 = ddddocr.DdddOcr(show_ad=False,
                       use_gpu=True)

ocr2 = ddddocr.DdddOcr(
    show_ad=False,
    use_gpu=True,
    beta=True)

# 验证码数据目录
img_dir = "您待打标签的验证码数据的根目录"
# 验证码数据集保存目录
save_dir = "您保存打好标签的验证码数据保存的根目录"
# 日志保存路径
log_dir = 'logs'  # 指定日志文件夹

""" ------ 日志配置无需更改 --------  """
# 如果文件夹不存在，则创建它
if not os.path.exists(log_dir):
    os.makedirs(log_dir)
if not os.path.exists(save_dir):
    os.makedirs(save_dir)
if not os.path.exists(img_dir):
    os.makedirs(img_dir)

log_file = os.path.join(log_dir, 'dataset.log')  # 完整的日志文件路径
# 配置日志的基本设置
logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S',
                    # 注意：这里我们不使用filename参数，因为我们将手动添加FileHandler
                    # filename=log_file,  # 如果使用basicConfig的filename，则不需要下面的FileHandler
                    filemode='a')  # 'a'表示追加模式，如果文件已存在则继续写入

# 创建一个FileHandler，用于写入日志文件
# 注意：由于我们已经在basicConfig中设置了基本的配置，这里只需要指定文件名
file_handler = logging.FileHandler(log_file)
file_handler.setLevel(logging.INFO)  # 可以为handler单独设置级别
logger = logging.getLogger()  # 或者使用 logging.getLogger(__name__) 来获取当前模块的logger
logger.addHandler(file_handler)  # 注意：如果你之前使用了basicConfig的filename参数，这里添加handler可能会导致日志被重复写入

""" ------- 功能函数 --------- """


# 添加一条日志
def add_log(txt):
    logger.info(txt)


# 通过自带模型和自己训练的模型进行双层识别
def check(img_path):
    with open(img_path, 'rb') as f:
        image = f.read()
    data1 = ocr2.classification(image, png_fix='True')
    data2 = ocr1.classification(image, png_fix='True')
    return data1, data2


# 两次识别结果不一样， 使用opencv展示图片
def show_err_img(name, img_path):
    img = cv2.imread(img_path)
    img = cv2.resize(img, (0, 0), fx=7, fy=7)
    cv2.imshow(name, img)
    # 等待键盘事件 输入Enter时关闭图片
    while True:
        # 等待键盘事件，参数为0表示无限期等待
        key = cv2.waitKey(0)
        # 检查是否按下了Enter键（在大多数环境中，Enter键会导致cv2.waitKey()返回-1或换行符的ASCII码，但-1更常见）
        if key == -1 or key == 13:  # 13是Enter键的ASCII码，但-1更常见
            break  # 退出循环
    # 关闭所有OpenCV窗口
    cv2.destroyAllWindows()


def selector(option1, option2):
    res = ''
    print("请选择一个选项:")
    print(f"p: 使用 [{option1}] 作为标签名字")
    print(f"[: 使用 [{option2}] 作为标签名字")
    print("]: 使用自定义选项作为标签名字")
    print("=: 删除图片")
    while True:
        choice = input("请输入选项编号:  'p' --- '[' ---- ']' ---- '+':\n 输入----> :")
        if choice == 'p':
            print("你选择了选项一。")
            res = option1
            break  # 如果不需要再次选择，可以退出循环
        elif choice == '[':
            print("你选择了选项二。")
            res = option2
            break  # 如果不需要再次选择，可以退出循环
        elif choice == ']':
            res = input('请输入标签名称:').strip()
            break
        elif choice == "=":
            break
        else:
            print("无效的输入，请重新输入。")
    os.system('cls')

    return res


# 循环遍历所有，比较两次识别结果是否相似
def make_dataset(path_dir, save_path):
    list_dir = os.listdir(path_dir)
    for img_name in list_dir:
        old_name = os.path.join(path_dir, img_name)
        result1, result2 = check(old_name)
        # 两次结果一样时
        if result1 == result2:
            new_name = os.path.join(save_path, result1 + '.' + img_name.split('.')[-1])
            add_log(f"图片: {new_name} -------- 识别一:{result1} --- 识别二:{result2}")
        # 两次结果不一致时
        else:
            show_err_img(img_name.split('.')[0], old_name)
            end_result = selector(result1, result2)
            # 如果最终结果是空的话就删除图片
            if end_result == '':
                os.remove(old_name)
                add_log(f'文件 {old_name} 已被删除。')
                continue
            new_name = os.path.join(save_path, end_result + '.' + img_name.split('.')[-1])
            add_log(f"图片: {new_name} -------- 识别一:{result1} --- 识别二:{result2} ----手选输入:{end_result}")
        os.rename(old_name, new_name)


if __name__ == '__main__':
    make_dataset(img_dir, save_dir)
